#include <mambo.hpp>
#include <mpi/communicator.h>
#include <mpi/collectives.h>
#include <dns_ops.h>
#include <cstdlib>
#include <iostream>
using namespace std;

int main(int argc, char** argv)
{

  mpi::communicator mycomm(argc,argv);

  typedef  double real_type;
  typedef  complex<double> complex_type;
  const int n_var=3;
  int niters=2;
  int nz=12288;
  int npx=512;
  int npy=512;
  bool debug=false;
  int ic=0;
  while(ic<argc)
  {
    string a(argv[ic]);
    if(a.find("-i")<a.size()) 
      niters=atoi(argv[++ic]);
    else if(a.find("-s")<a.size()) //get the z
      nz=atoi(argv[++ic]);
    else if(a.find("-npx")<a.size()) //number of pencils
      npx=atoi(argv[++ic]);
    else if(a.find("-npy")<a.size()) //number of pencils
      npy=atoi(argv[++ic]);
    else if(a.find("-debug")<a.size())  
      debug=true;
    else if(a.find("-h")<a.size())
    {
      cerr << "Usage:\n  rk4 -s fft-size -n fft-bundles-per-thread -i iterations" << endl;
      cerr << "Defaults: \n rk4 -s 12288 -n 512 -i 2 " << endl;
    }
    ++ic;
  }

  mycomm.log() << "BW nx_loc=" << nz/npx << " ny_loc=" << nz/npy << " nz_loc=" << nz
    << " iterations = " << niters << endl;

  int nc=nz/npx*nz/npy;
  int ntot=n_var*nz*nc;
  int np=omp_get_max_threads();

  vector<real_type*> Q(np), rhs(np),Q_tmp(np),Q_old(np);
#pragma omp parallel
  {
    int ip=omp_get_thread_num();
    Q[ip]=new real_type[ntot];
    rhs[ip]=new real_type[ntot];
    Q_tmp[ip]=new real_type[ntot];
    Q_old[ip]=new real_type[ntot];
    std::fill(Q[ip],Q[ip]+ntot,0.0);
    std::fill(rhs[ip],rhs[ip]+ntot,0.0);
    std::fill(Q_tmp[ip],Q_tmp[ip]+ntot,0.0);
    std::fill(Q_old[ip],Q_old[ip]+ntot,0.0);
  }

  mycomm.barrier();
  if(!mycomm.rank()) cout << "PHASEC" << endl;
  Timer wtimer;
  vector<double> t(omp_get_max_threads(),0.0);
#pragma omp parallel 
  {
    int ip=omp_get_thread_num();
    const double one_third=1.0/3.0;
    const double delta=0.01/2.0;
    double tot_loc=0.0;
    //Q +=  one_third*rhs
    //Q_tmp=Q_old+delta*rhs
    Timer myclock;
    char s[128];
    for(int i=0; i<niters; ++i)
    {
      myclock.restart();
#if defined(RK4_CLOOP)
      dns_ops::apply(Q[ip],Q_tmp[ip],Q_old[ip],rhs[ip],one_third,delta,ntot);
      dns_ops::apply(rhs[ip],Q_tmp[ip],Q_old[ip],Q[ip],one_third,delta,ntot);
#endif
#if defined(RK4_FLOOP)
      do_rk4(Q[ip],Q_tmp[ip],Q_old[ip],rhs[ip],&one_third,&delta,&ntot);
      do_rk4(rhs[ip],Q_tmp[ip],Q_old[ip],Q[ip],&one_third,&delta,&ntot);
#endif
#if defined(RK4_2LOOPS)
      dns_ops::apply1(Q[ip],rhs[ip],one_third,ntot);
      dns_ops::apply2(Q_tmp[ip],Q_old[ip],rhs[ip],delta,ntot);
      //double tt=myclock.elapsed();
      //printf("BW LOOP=%i PE=%i FIRST %12.5e\n",i,ip,tt);
      dns_ops::apply1(rhs[ip],Q[ip],one_third,ntot);
      dns_ops::apply2(Q_tmp[ip],Q_old[ip],Q[ip],delta,ntot);
#endif
      tot_loc+=myclock.elapsed();
      //printf("BW LOOP=%i PE=%i %12.5e %12.5e\n",i,ip,tot_loc/static_cast<double>(i+1),tot_loc);
    }
    t[ip]=tot_loc;
  }
  double tot_time=wtimer.elapsed();
  mycomm.barrier();
  if(!mycomm.rank()) cout << "PHASED" << endl;

  for(int ip=1;ip<np;++ip) t[0]+=t[ip];

  //vector<double> dt(4);
  //dt[0]=t0;dt[1]=t1;dt[2]=t2;dt[3]=tot_time;
  //mpi::reduce(mycomm,dt);

  cout.setf(ios::fixed, ios::floatfield);
  cout.precision(5);
  cout  << "BW OMP_NUM_THREADS = " << omp_get_max_threads() 
    <<  " Wall-clock (sec) = " << tot_time 
    << " sum = " << t[0]/static_cast<double>(omp_get_max_threads())
    << endl;
 
  cout << "BW fft-dim num-ffts ops 2xRK   2xRK-wc" << endl;
  cout << "BW " << nz << " " << nc  << " " << nz*nc*n_var
       << " " << t[0]/static_cast<double>(niters*omp_get_max_threads())
       << " " << tot_time/static_cast<double>(niters)
       << endl;
  //const double rk4=4.0;
  //double nodefac= rk4/static_cast<double>(niters*omp_get_max_threads());
  //cout << "BW fft-dim num-ffts ops c-loop f-loop c-2loops" << endl;
  //cout << "BW " << nz << " " << nc  << " " << nz*nc*n_var
  //  << " " << t[0]*nodefac << " " << t[1]*nodefac  << " " << t[2]*nodefac
  //  << endl;
  for(int ip=0;ip<np; ++ip) delete Q_old[ip];
  for(int ip=0;ip<np; ++ip) delete Q_tmp[ip];
  for(int ip=0;ip<np; ++ip) delete rhs[ip];
  for(int ip=0;ip<np; ++ip) delete Q[ip];
  return 0;
}

